package padding

import (
	"bytes"
	"crypto/aes"
	"errors"
)

//Padding ...
type Padding interface {
	Padding(src []byte) []byte
	UnPadding(src []byte) ([]byte, error)
}

//CbcPatddingPkcs7 CBC_PADDING_PKCS7
func CbcPatddingPkcs7() Padding {
	return &cbcPaddingPkcs7{}
}

type cbcPaddingPkcs7 struct {
}

//Padding ...
func (p *cbcPaddingPkcs7) Padding(src []byte) []byte {
	return pkcs7Padding(src)
}

//UnPadding ...
func (p *cbcPaddingPkcs7) UnPadding(src []byte) ([]byte, error) {
	return pkcs7UnPadding(src)
}
func pkcs7Padding(src []byte) []byte {
	padding := aes.BlockSize - len(src)%aes.BlockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(src, padtext...)
}

func pkcs7UnPadding(src []byte) ([]byte, error) {
	length := len(src)
	unpadding := int(src[length-1])

	if unpadding > aes.BlockSize || unpadding == 0 {
		return nil, errors.New("Invalid pkcs7 padding (unpadding > aes.BlockSize || unpadding == 0)")
	}

	pad := src[len(src)-unpadding:]
	for i := 0; i < unpadding; i++ {
		if pad[i] != byte(unpadding) {
			return nil, errors.New("Invalid pkcs7 padding (pad[i] != unpadding)")
		}
	}

	return src[:(length - unpadding)], nil
}
